package bigo.lib;

import bigo.BigODriver;
import bigo.data.RowIndex;
import bigo.data.Vector;
import bigo.data.Vector.Wrapper;
import java.io.IOException;
import java.util.HashMap;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.Mapper.Context;
import org.apache.hadoop.mapreduce.Reducer;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;

/**
 *
 * @author songliu
 */
public class AgreegationDriver extends BigODriver {

    @Mapper
    public static class AgreegationMapper extends org.apache.hadoop.mapreduce.Mapper<IntWritable, Wrapper, RowIndex, Wrapper> {

        @Override
        protected void map(IntWritable key, Wrapper value, Context context)
                throws IOException, InterruptedException {
            FileSplit s = (FileSplit) context.getInputSplit();
            String source = s.getPath().getParent().getName();
            context.write(new RowIndex(key.get(), source), value);
        }

        @Override
        protected void setup(Context context) throws IOException, InterruptedException {
            super.setup(context);
            BigODriver.loadFields(context.getConfiguration());
        }
    }

    @Combiner
    public static abstract class AgreegationCombiner extends org.apache.hadoop.mapreduce.Reducer<RowIndex, Wrapper, RowIndex, Wrapper> {

        public abstract Vector agreegate(RowIndex index, HashMap<String, Vector> vectors);

        @Override
        protected void reduce(RowIndex key, Iterable<Wrapper> values, Context context) throws IOException, InterruptedException {
            HashMap<String, Vector> vectors = new HashMap<String, Vector>();
            for (Wrapper w : values) {
                Vector v = (Vector) w.get();
                vectors.put(key.sourceMatrix, v);
            }

            context.write(key,
                    new Wrapper(agreegate(key, vectors)));
        }

        @Override
        protected void setup(Reducer.Context context) throws IOException, InterruptedException {
            super.setup(context);
            BigODriver.loadFields(context.getConfiguration());
        }
    }

    @bigo.lib.Reducer
    public static abstract class AgreegationReducer extends org.apache.hadoop.mapreduce.Reducer<RowIndex, Wrapper, IntWritable, Wrapper> {

        public abstract Vector agreegate(RowIndex index, HashMap<String, Vector> vectors);

        @Override
        protected void reduce(RowIndex key, Iterable<Wrapper> values, Context context) throws IOException, InterruptedException {
            HashMap<String, Vector> vectors = new HashMap<String, Vector>();
            for (Wrapper w : values) {
                Vector v = (Vector) w.get();
                vectors.put(key.sourceMatrix, v);
            }

            context.write(new IntWritable(key.rowIndex),
                    new Wrapper(agreegate(key, vectors)));
        }

        @Override
        protected void setup(Reducer.Context context) throws IOException, InterruptedException {
            super.setup(context);
            BigODriver.loadFields(context.getConfiguration());
        }
    }

    @Override
    public void init(Job job) {
        job.setMapOutputKeyClass(RowIndex.class);
        job.setSortComparatorClass(RowIndex.Comparator.class);
        job.setGroupingComparatorClass(RowIndex.GroupComparator.class);
    }

    @Override
    public void preRun(Job job) {
    }
}
